#pragma once

#include "config.hpp"
#include "logical_type.hpp"
#include "rc.hpp"
#include <arrow/api.h>
#include <iostream>
#include <memory>
#include <mutex>
#include <numa.h>
#include <sched.h>
#include <string>
#include <thread>
#include <unordered_map>
#include <vector>

namespace DaseX {

class ColumnChunk {
public:
    std::shared_ptr<arrow::Array> data;

public:
    ColumnChunk() = default;

    ColumnChunk(std::shared_ptr<arrow::Array> data_) : data(std::move(data_)) {};

    arrow::Status append_int_array(std::vector<int> &array) {
        arrow::Int32Builder int32builder;
        ARROW_RETURN_NOT_OK(int32builder.AppendValues(array));
        ARROW_ASSIGN_OR_RAISE(data, int32builder.Finish());
        return arrow::Status::OK();
    }

    RC append_int_data(std::vector<int> &array) {
        arrow::Status st = append_int_array(array);
        RC rc = RC::FAILED;
        if (st.ok()) {
            rc = RC::SUCCESS;
        }
        return rc;
    }

    arrow::Status append_float_array(std::vector<float> &array);

    arrow::Status append_float_data(std::vector<float> &array);

    arrow::Status append_double_array(std::vector<double> &array);

    arrow::Status append_double_data(std::vector<double> &array);

    arrow::Status append_string_array(std::vector<std::string> &array) {
        arrow::StringBuilder stringbuilder;
        ARROW_RETURN_NOT_OK(stringbuilder.AppendValues(array));
        ARROW_ASSIGN_OR_RAISE(data, stringbuilder.Finish());
        return arrow::Status::OK();
    }

    RC append_string_data(std::vector<std::string> &array) {
        arrow::Status st = append_string_array(array);
        RC rc = RC::FAILED;
        if (st.ok()) {
            rc = RC::SUCCESS;
        }
        return rc;
    }

    int get_column_chunk_size() const { return data->length(); }

    bool is_null(int64_t i) const { return data->IsNull(i); }

    std::shared_ptr<arrow::Array> get_data() { return data; }
};

class Column {
public:
    std::string colunm_name;
    std::vector<std::shared_ptr<ColumnChunk>> chunks;
    int chunk_nums = 0;
    std::mutex mutex_colunm;

public:
    Column(std::string colunm_name_) : colunm_name(colunm_name_) {}

    void insert_data(std::shared_ptr<ColumnChunk> chunk) {
        std::lock_guard<std::mutex> lock(mutex_colunm);
        chunks.push_back(chunk);
        chunk_nums++;
    }

    std::shared_ptr<ColumnChunk> get_data(int idx) const {
        if (idx > chunk_nums) {
            throw std::runtime_error("Out of Range");
        }
        return chunks[idx];
    }

    int get_chunk_nums() const { return chunk_nums; };
};

class Field {
public:
    std::string field_name;
    LogicalType type;
    std::shared_ptr<arrow::Field> field;

public:
    Field(std::string field_name_, LogicalType type_)
            : field_name(field_name_), type(type_) {
        switch (type) {
            case LogicalType::INTEGER:
                field = arrow::field(field_name, arrow::int32());
                break;
            case LogicalType::FLOAT:
                field = arrow::field(field_name, arrow::float32());
                break;
            case LogicalType::DOUBLE:
                field = arrow::field(field_name, arrow::float64());
                break;
            case LogicalType::STRING:
                field = arrow::field(field_name, arrow::utf8());
                break;
            default:
                break;
        }
    }
};

class TableInfo {
public:
    std::string table_name;
    std::vector<std::string> filed_names;
    std::vector<LogicalType> filed_types;
    int column_nums = 0;
    int partition_num = 0;

public:
    TableInfo() = default;

    TableInfo(const std::string &table_name_,
              std::vector<std::string> &filed_names_,
              std::vector<LogicalType> &filed_types_,
              int partition_num_ = Numa::CORES)
            : table_name(table_name_), filed_names(filed_names_),
              filed_types(filed_types_), partition_num(partition_num_) {
        column_nums = filed_names.size();
    }
};

class PartitionedTable {
public:
    std::vector<std::shared_ptr<Field>> fileds;
    std::shared_ptr<arrow::Schema> schema;
    std::vector<std::shared_ptr<Column>> columns;
    int chunk_nums = 0;  // 列块
    int column_nums = 0; // 列数
    int row_nums = 0;    // 分区表共有多少行数据
    int partition_id;
    std::mutex mutex_table;

public:
    PartitionedTable(int partition_id_) : partition_id(partition_id_) {}

    /// 从元数据（CataLog）中获取TableInfo,用从元数据中获取TableInfo初始化fileds +
    /// schema + columns
    void initailize(const TableInfo &table_info) {
        column_nums = table_info.column_nums;
        if (column_nums == 0) {
            throw std::runtime_error("Table is empty!!!");
        }
        arrow::FieldVector filed_vector;
        for (int i = 0; i < column_nums; i++) {
            fileds.push_back(std::move(std::make_shared<Field>(
                    table_info.filed_names[i], table_info.filed_types[i])));
            filed_vector.push_back(fileds[i]->field);
            columns.push_back(std::make_shared<Column>(table_info.filed_names[i]));
        }
        schema = arrow::schema(filed_vector);
    } // initailize_table

    std::shared_ptr<arrow::RecordBatch>
    get_data(int idx, const std::vector<int> &project_ids = std::vector<int>()) {
        std::shared_ptr<arrow::RecordBatch> rbatch;
        if (idx >= chunk_nums) {
            arrow::FieldVector res_fields_vector;
            for(int i = 0; i < project_ids.size(); i++) {
                res_fields_vector.push_back(schema->field(project_ids[i]));
            }
            std::shared_ptr<arrow::Schema> res_schema = std::make_shared<arrow::Schema>(res_fields_vector);
            rbatch = arrow::RecordBatch::MakeEmpty(res_schema).ValueOrDie();
            return rbatch;
        }
        arrow::ArrayVector column_vector;
        int length = 0;
        for (int i = 0; i < columns.size(); i++) {
            std::shared_ptr<ColumnChunk> column_chunk = columns[i]->get_data(idx);
            column_vector.push_back(column_chunk->data);
        }
        length = column_vector[0]->length();
        rbatch = arrow::RecordBatch::Make(schema, length, column_vector);
        if (project_ids.size() == 0) {
            return rbatch->Slice(0);
        }
        std::shared_ptr<arrow::RecordBatch> tmp =
                rbatch->SelectColumns(project_ids).ValueOrDie();
        return tmp->Slice(0);
    }

    RC insert_data(std::shared_ptr<arrow::RecordBatch> &record_batch) {
        RC rc;
        if (record_batch->num_rows() == 0 || record_batch->num_columns() == 0) {
            rc = RC::INVALID_DATA;
            return rc;
        }
        const std::vector<std::shared_ptr<arrow::Array>> &column_vector =
                record_batch->columns();
        for (int i = 0; i < column_vector.size(); i++) {
            std::lock_guard<std::mutex> lock(mutex_table);
            columns[i]->insert_data(
                    std::move(std::make_shared<ColumnChunk>(column_vector[i])));
        }
        row_nums += record_batch->num_rows();
        chunk_nums++;
        return rc;
    } // insert_data

    /// 获取分区所在的numa节点
    int get_socket() { return Numa::get_socket(partition_id); }

    /// FOR DEBUG
    void print_table() {
        std::cout << "==============================partition" << partition_id
                  << std::endl;
        int chunk_nums = columns[0]->chunk_nums;

        for (int i = 0; i < chunk_nums; i++) {
            std::shared_ptr<arrow::RecordBatch> rbatch = get_data(i);
            std::cout << "=======================chunk_nums" << i << std::endl;
            std::cout << rbatch->ToString() << std::endl;
        }
    }
};

class Table {
public:
    std::string db = "db_314";
    std::string table_name;
    std::vector<std::shared_ptr<Field>> fileds;
    std::shared_ptr<arrow::Schema> schema;
    std::vector<std::shared_ptr<PartitionedTable>> partition_tables;

    int partition_num;                 // 逻辑分区数
    int numa_num;                      // numa节点的个数
    int column_nums = 0;               // 列数
    int row_nums = 0;                  // 表共有多少行数据
    std::vector<int> partition_bitmap; // 分区有数据为1，没有数据为0
public:
    Table(TableInfo &table_info_) { initailize(table_info_); }

    void initailize(TableInfo &table_info) {
        table_name = table_info.table_name;
        partition_num = table_info.partition_num;
        for (int i = 0; i < partition_num; ++i) {
            partition_tables.push_back(std::make_shared<PartitionedTable>(i));
            partition_tables[i]->initailize(table_info);
            partition_bitmap.push_back(0);
        }

        arrow::FieldVector filed_vector;
        column_nums = table_info.column_nums;
        for (int i = 0; i < column_nums; i++) {
            fileds.push_back(std::move(std::make_shared<Field>(
                    table_info.filed_names[i], table_info.filed_types[i])));
            filed_vector.push_back(fileds[i]->field);
        }
        schema = arrow::schema(filed_vector);
    } // initailize_table

    RC insert_data(std::shared_ptr<arrow::RecordBatch> record_batch,
                   int partition_id) {
        RC rc = RC::FAILED;
        if (record_batch->schema()->Equals(schema)) {
            if (partition_id >= partition_num) {
                return rc;
            }
            if (record_batch->num_rows() == 0 || record_batch->num_columns() == 0) {
                rc = RC::INVALID_DATA;
                return rc;
            }

            partition_tables[partition_id]->insert_data(record_batch);
            partition_bitmap.at(partition_id) = 1;
            row_nums += record_batch->num_rows();
            rc = RC::SUCCESS;
        }
        return rc;
    } // insert_data

    void print_table() {
        for (int i = 0; i < partition_num; ++i) {
            if (partition_bitmap[i]) {
                partition_tables[i]->print_table();
            }
        }
    }
};

class CataLog {
public:
    std::unordered_map<std::string, TableInfo> table_infos;
    std::unordered_map<std::string, std::shared_ptr<Table>> tables;
    std::mutex mutex_catalog;

public:
    CataLog() {}

    RC register_table(TableInfo &table_info);

    RC delete_table(std::string table_name);

    RC get_table(std::string table_name, std::shared_ptr<Table> &table);
};

extern CataLog global_catalog;

} // namespace DaseX